import pygame
import math


class CartPoleVisualizer:
    def __init__(
        self, width=1000, height=800, cart_width=50, cart_height=20, pole_length=1.5
    ):
        pygame.init()
        self.width = width
        self.height = height
        self.screen = pygame.display.set_mode((self.width, self.height))
        pygame.display.set_caption("3D CartPole Visualization")

        self.cart_width = cart_width
        self.cart_height = cart_height
        self.pole_length = pole_length
        self.scale_x = self.width / 5
        self.scale_y = self.width / 10
        self.scale_z = self.height / 4

        self.cart_y = self.height // 2
        self.center_x = self.width // 2

        self.WHITE = (255, 255, 255)
        self.BLACK = (0, 0, 0)
        self.BLUE = (0, 0, 255)
        self.RED = (255, 0, 0)

        self.running = True

    def reset_position(self):
        self.x = 0
        self.y = 0
        self.z = 0
        self.pole_angle = 0.0

    def world_to_screen(self, x, y, z):
        screen_x = self.center_x + int(x * self.scale_x)
        screen_y = self.cart_y - int(z * self.scale_z)
        return screen_x, screen_y

    def draw_cartpole(self, x, y, z, pole_angle):
        self.screen.fill(self.WHITE)

        cart_x, cart_z = self.world_to_screen(x, y, z)

        pygame.draw.line(self.screen, self.RED, (0, cart_z), (self.width, cart_z), 2)
        pygame.draw.rect(
            self.screen,
            self.BLUE,
            (
                cart_x - self.cart_width // 2,
                cart_z - self.cart_height // 2,
                self.cart_width,
                self.cart_height,
            ),
        )

        pole_x = x + self.pole_length * math.sin(pole_angle)
        pole_z = z - self.pole_length * math.cos(pole_angle)
        pole_screen_x, pole_screen_z = self.world_to_screen(pole_x, y, pole_z)

        pygame.draw.line(
            self.screen, self.BLACK, (cart_x, cart_z), (pole_screen_x, pole_screen_z), 6
        )
        # self.draw_reset_button()

        pygame.display.flip()

    def draw_reset_button(self):
        button_rect = pygame.Rect(self.width - 120, 20, 100, 40)
        pygame.draw.rect(self.screen, self.RED, button_rect)
        font = pygame.font.Font(None, 24)
        text = font.render("Reset", True, self.WHITE)
        text_rect = text.get_rect(center=button_rect.center)
        self.screen.blit(text, text_rect)

    def is_reset_button_clicked(self, mouse_pos):
        return (
            self.width - 120 <= mouse_pos[0] <= self.width - 20
            and 20 <= mouse_pos[1] <= 60
        )

    def update(self, x, y, z, pole_angle):
        self.draw_cartpole(x, y, z, pole_angle)

    def quit(self):
        self.running = False
        pygame.quit()


if __name__ == "__main__":
    visualizer = CartPoleVisualizer()

    x, y, z = 0, 0, 0
    pole_angle = math.pi / 6

    while visualizer.running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                visualizer.quit()

        visualizer.update(x, y, z, pole_angle)
